package com.example.llm;

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;
import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.function.Consumer;
import java.util.ArrayList;
import java.util.List;

/**
 * Ollama API客户端类
 * 负责与Ollama服务器进行通信，发送请求并处理响应
 */
public class OllamaClient {
    private final String baseUrl;        // Ollama服务器的基础URL
    private final Gson gson;             // 用于JSON序列化和反序列化的Gson实例
    private Consumer<String> debugCallback;  // 调试信息的回调函数
    
    /**
     * 构造函数
     * @param baseUrl Ollama服务器的地址
     */
    public OllamaClient(String baseUrl) {
        this.baseUrl = baseUrl;
        this.gson = new Gson(); // 初始化Gson实例
    }
    
    /**
     * 设置调试回调函数
     * @param callback 用于接收调试信息的回调函数
     */
    public void setDebugCallback(Consumer<String> callback) {
        this.debugCallback = callback;
    }
    
    /**
     * 输出调试信息
     * @param message 调试信息内容
     */
    private void debug(String message) {
        if (debugCallback != null) {
            debugCallback.accept(message); // 调用回调函数输出调试信息
        }
    }
    
    /**
     * 创建新的提示构建器
     * @return PromptBuilder实例，用于构建API请求
     */
    public PromptBuilder prompt() {
        try {
            return new PromptBuilder();
        } catch (IOException e) {
            throw new RuntimeException("Failed to create PromptBuilder", e);
        }
    }
    
    /**
     * 提示构建器内部类
     * 使用Builder模式构建API请求
     */
    public class PromptBuilder {
        private String model = "deepseek:latest";  // 默认模型
        private double temperature = 0.7;          // 温度参数
        private final List<Message> history = new ArrayList<>();  // 历史消息
        private String userMessage;                // 用户消息
        private final StringBuilder systemSpec = new StringBuilder();  // 系统提示词
        private final URL url;                     // API请求URL
        private int maxContext = 4096;             // 最大上下文长度
        
        /**
         * 构造函数
         * @throws IOException 如果URL创建失败
         */
        public PromptBuilder() throws IOException {
            this.url = new URL(baseUrl + "/api/generate"); // 设置请求URL
        }
        
        /**
         * 设置模型
         * @param model 模型名称
         * @return 当前PromptBuilder实例
         */
        public PromptBuilder withModel(String model) {
            this.model = model;
            return this;
        }
        
        /**
         * 设置温度参数
         * @param temperature 温度值
         * @return 当前PromptBuilder实例
         */
        public PromptBuilder withTemperature(double temperature) {
            this.temperature = temperature;
            return this;
        }
        
        /**
         * 设置系统消息
         * @param message 系统消息内容
         * @return 当前PromptBuilder实例
         */
        public PromptBuilder withSystemMessage(String message) {
            systemSpec.append(message);
            return this;
        }
        
        /**
         * 设置历史消息
         * @param history 历史消息列表
         * @return 当前PromptBuilder实例
         */
        public PromptBuilder withHistory(List<Message> history) {
            this.history.clear();
            this.history.addAll(history);
            return this;
        }
        
        /**
         * 设置用户消息
         * @param message 用户输入的消息
         * @return 当前PromptBuilder实例
         */
        public PromptBuilder withUserMessage(String message) {
            this.userMessage = message;
            return this;
        }
        
        /**
         * 设置最大上下文长度
         * @param maxContext 最大上下文长度
         * @return 当前PromptBuilder实例
         */
        public PromptBuilder withMaxContext(int maxContext) {
            this.maxContext = maxContext;
            return this;
        }
        
        /**
         * 设置系统参数
         * @param specBuilder 系统参数构建器
         * @return 当前PromptBuilder实例
         */
        public PromptBuilder withSystem(Consumer<SystemSpec> specBuilder) {
            SystemSpec spec = new SystemSpec();
            specBuilder.accept(spec);
            systemSpec.append(spec.toString());
            return this;
        }
        
        /**
         * 发送流式请求
         * @return Flux<String> 响应内容的流
         * @throws IOException 如果请求失败
         */
        public Flux<String> stream() throws IOException {
            return Flux.create(onNext -> {
                try {
                    HttpURLConnection connection = createConnection(); // 创建HTTP连接
                    String requestBody = gson.toJson(build()); // 构建请求体
                    
                    // 发送请求
                    try (OutputStream os = connection.getOutputStream()) {
                        byte[] input = requestBody.getBytes(StandardCharsets.UTF_8);
                        os.write(input, 0, input.length);
                    }
                    
                    // 检查响应状态
                    if (connection.getResponseCode() != HttpURLConnection.HTTP_OK) {
                        throw new IOException("HTTP error code: " + connection.getResponseCode());
                    }
                    
                    processStream(onNext, connection); // 处理响应流
                } catch (Exception e) {
                    if (debugCallback != null) {
                        debugCallback.accept("处理响应时发生错误: " + e.getMessage());
                    }
                    onNext.error(e); // 发送错误
                }
            });
        }
        
        /**
         * 创建HTTP连接
         * @return HttpURLConnection 实例
         * @throws IOException 如果连接失败
         */
        private HttpURLConnection createConnection() throws IOException {
            HttpURLConnection conn = (HttpURLConnection) url.openConnection();
            conn.setRequestMethod("POST");
            conn.setRequestProperty("Content-Type", "application/json");
            conn.setDoOutput(true);
            return conn;
        }
        
        /**
         * 构建请求JSON对象
         * @return JsonObject 包含请求参数的JSON对象
         */
        private JsonObject build() {
            JsonObject json = new JsonObject();
            json.addProperty("model", model);
            json.addProperty("temperature", temperature);
            json.addProperty("context_length", maxContext);
            
            StringBuilder prompt = new StringBuilder();
            
            // 添加系统提示词
            if (!systemSpec.toString().isEmpty()) {
                prompt.append(systemSpec.toString()).append("\n\n");
            }
            
            // 添加历史消息
            for (Message msg : history) {
                prompt.append(msg.role.equals("user") ? "Human: " : "Assistant: ")
                      .append(msg.content)
                      .append("\n\n");
            }
            
            // 添加当前用户消息
            prompt.append("Human: ").append(userMessage).append("\n\nAssistant: ");
            
            json.addProperty("prompt", prompt.toString());
            json.addProperty("stream", true);
            
            debug("发送请求到: " + url);
            debug("请求内容: " + json);
            
            return json;
        }
    }
    
    /**
     * 处理流式响应
     * @param onNext Flux发射器
     * @param connection HTTP连接
     * @throws IOException 如果处理失败
     */
    private void processStream(FluxSink<String> onNext, HttpURLConnection connection) throws IOException {
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(connection.getInputStream(), StandardCharsets.UTF_8))) {
            String line;
            StringBuilder currentResponse = new StringBuilder();
            boolean isDone = false;
            int retryCount = 0;
            final int MAX_RETRIES = 3;              // 最大重试次数
            final int BUFFER_SIZE = 8192;           // 读取缓冲区大小
            
            char[] buffer = new char[BUFFER_SIZE];
            int bytesRead;
            
            // 读取响应流
            while ((bytesRead = reader.read(buffer)) != -1) {
                String chunk = new String(buffer, 0, bytesRead);
                String[] lines = chunk.split("\n");
                
                for (String currentLine : lines) {
                    if (currentLine.trim().isEmpty()) {
                        continue; // 跳过空行
                    }
                    
                    try {
                        // 解析JSON响应
                        JsonObject response = gson.fromJson(currentLine, JsonObject.class);
                        String content = response.get("response").getAsString();
                        boolean done = response.get("done").getAsBoolean();
                        
                        // 输出调试信息
                        if (debugCallback != null) {
                            debugCallback.accept("收到响应内容: " + content);
                            debugCallback.accept("响应完成状态: " + done);
                        }
                        
                        // 处理响应内容
                        if (!content.isEmpty()) {
                            currentResponse.append(content);
                            onNext.next(content); // 发送内容
                        }
                        
                        // 检查是否完成
                        if (done) {
                            isDone = true;
                            onNext.complete(); // 完成流
                            return;
                        }
                        
                        retryCount = 0; // 重置重试计数
                    } catch (Exception e) {
                        // 错误重试处理
                        if (retryCount < MAX_RETRIES) {
                            retryCount++;
                            debugCallback.accept("解析响应时出错，重试 " + retryCount + "/" + MAX_RETRIES);
                            continue;
                        }
                        throw new RuntimeException("解析响应失败: " + e.getMessage());
                    }
                }
            }
            
            // 检查是否正常完成
            if (!isDone) {
                onNext.error(new RuntimeException("Stream ended without completion"));
            }
        } catch (Exception e) {
            onNext.error(e); // 发送错误
        }
    }
    
    /**
     * 系统提示词规范类
     */
    public static class SystemSpec {
        private final StringBuilder params = new StringBuilder();
        
        public SystemSpec param(String key, String value) {
            if (params.length() > 0) {
                params.append("\n");
            }
            params.append(key).append(": ").append(value);
            return this;
        }
        
        @Override
        public String toString() {
            return params.toString(); // 返回系统参数字符串
        }
    }
    
    /**
     * 请求顾问接口
     */
    public interface Advisor {
        void apply(JsonObject json); // 应用顾问逻辑
    }
    
    /**
     * 消息类
     * 用于存储对话消息
     */
    public static class Message {
        private final String role;      // 角色（用户/助手）
        private final String content;   // 消息内容
        
        public Message(String role, String content) {
            this.role = role;
            this.content = content;
        }
        
        public String getRole() {
            return role; // 获取角色
        }
        
        public String getContent() {
            return content; // 获取内容
        }
    }
}